import streamlit as st
import os
import requests
import base64
from colivara_py import ColiVara
from pathlib import Path

# 获取脚本当前目录
current_dir = os.path.dirname(__file__)

# 指定保存目录为当前目录下的 docs 子目录
save_dir = os.path.join(current_dir, "docs")

# 创建保存目录（如果不存在）
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

# ColiVara 服务器 API URL
COLIVARA_BASE_URL = "http://localhost:8001"
COLIVARA_API_KEY = "你启动服务产生的Token"
#COLIVARA_BASE_URL = "https://api.colivara.com"
#COLIVARA_API_KEY = "你的API-KEY"

rag_client = ColiVara(base_url=COLIVARA_BASE_URL, api_key=COLIVARA_API_KEY)

def sync_and_delete_documents():
    """同步 docs 目录中的**所有**文件到 ColiVara；
       只有全部成功后才删除本地文件"""
    docs_path = Path(save_dir)
    files = [f for f in docs_path.glob('*') if f.is_file()]
    if not files:
        st.info("docs 目录为空，无需同步。")
        return

    # 收集同步结果
    failed_files = []
    for file in files:
        with open(file, 'rb') as f:
            encoded = base64.b64encode(f.read()).decode()
        try:
            rag_client.upsert_document(
                name=file.name,
                document_base64=encoded,
                collection_name="这个名字可以自定义",
                wait=True
            )
            st.success(f"✅ 已解析：{file.name}")
        except Exception as e:
            st.error(f"❌ 解析失败：{file.name} - {e}")
            failed_files.append(file.name)

    # 全部成功后再删除
    if not failed_files:
        for f in files:
            f.unlink()
        st.success("所有文件解析完成，本地docs目录文件已删除。")
    else:
        st.warning(f"以下文件未解析成功，本地文件未删除：{', '.join(failed_files)}")

st.title("**ColiVara 文档上传和嵌入应用**")

# 文件上传组件
uploaded_files = st.file_uploader("选择任意格式文档、图片文件上传",accept_multiple_files=True)

if uploaded_files:
    for uploaded_file in uploaded_files:
        file_name = uploaded_file.name
        file_path = os.path.join(save_dir, file_name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getvalue())
    st.success(f"已成功上传 {len(uploaded_files)} 个文件到 docs 目录")
        #st.success("文件已成功上传docs目录")

         # 添加一个按钮，用于手动同步和删除 docs 目录中的所有文件
    if st.button("文件已成功上传，点击解析后删除docs目录中的所有文件", key="sync_all"):
        sync_and_delete_documents()
